import cv2
import numpy as np
from ultralytics import YOLO  # 导入Ultralytics的YOLOv8模型库
from typing import Tuple, List, Dict, Any, Optional
from pathlib import Path
import torch  # 用于判断CUDA（GPU）是否可用
import os


# 定义yolov8检测类
class YOLOv8Detector:
    # 创建初始化方法，当创建类实例时会自动调用
    def __init__(self, config_loader):
        self.config = config_loader.get_model_config()      # 获取模型配置
        self.app_config = config_loader.get_app_config()    # 获取应用配置
        self.model_person, self.model_fire = None, None     # 初始化模型
        self.is_initialized = False                         # 初始化标志
        self._load_model()                                  # 加载模型
        self.dispose = False

    # 模型加载方法(内部使用)
    def _load_model(self):
        try:
            device = 'cpu' if self.config['model']['person']['device'] == 'cuda' and not torch.cuda.is_available() else self.config['model']['person']['device']  # 判断是否使用GPU
            print(f"当前使用设备为: {device}")                                     # 打印当前使用设备

            self.model_person = YOLO(self.config['model']['person']['path'])    # 加载人员检测模型
            self.model_person.to(device)                                        # 将模型移动到指定设备

            self.model_fire = YOLO(self.config['model']['fire']['path'])        # 加载火焰检测模型
            self.model_fire.to(device)                                          # 将模型移动到指定设备

            self.is_initialized = True                                          # 设置初始化标志
            print(
                f"模型加载成功: person--{self.config['model']['person']['path']}  "
                f"fire--{self.config['model']['fire']['path']}")                # 打印模型加载成功信息
        except Exception as e:
            print(f"模型加载失败: {e}")                                            # 打印模型加载失败信息
            self.is_initialized = False

    # 处理检测结果的内部方法(多个模型复用)
    def _process_detections(self, model, frame, model_cfg, target_classes):     # 处理检测结果
        detections = []                                                         # 初始化检测结果列表
        has_danger = False                                                      # 初始化危险标志
        results = model(frame, imgsz=model_cfg['imgsz'],                        # 进行检测
                        conf=self.app_config['detection']['confidence_threshold'],
                        max_det=self.config['optimization']['max_det'],
                        verbose=False)

        for r in results:
            if r.boxes is None: continue                # 如果检测结果为空，则跳过
            for box in r.boxes:
                cls_id = int(box.cls)                   # 获取类别ID
                conf = float(box.conf)                  # 获取置信度
                xyxy = box.xyxy[0].tolist()             # 获取边界框坐标

                for cls_name, cfg in target_classes.items():        # 遍历目标类别
                    if cls_id == cfg['id']:                         # 如果类别ID与目标类别ID匹配
                        detections.append({                         # 添加检测结果
                            'class': cls_name,                      # 类别名称
                            'confidence': conf,                     # 置信度
                            'bbox': xyxy,                           # 边界框坐标
                            'color': cfg['color'],                  # 颜色
                            'is_danger': cfg['is_danger']           # 是否危险
                        })
                        if cfg['is_danger']:                        # 如果类别为危险类别
                            has_danger = True
        return detections, has_danger                               # 返回检测结果和危险标志

    # 主要的检测方法(外部调用)
    def detect(self, frame: np.ndarray) -> Tuple[bool, List[Dict[str, Any]], bool]:
        if not self.is_initialized:                             # 如果模型未初始化
            return False, [], False                             # 返回False, 空列表和False

        try:
            detections = []                                     # 初始化检测结果列表
            has_danger = False                                  # 初始化危险标志
            detection_config = self.app_config['detection']['classes']      # 获取检测配置

            # 处理人员检测
            person_dets, person_danger = self._process_detections(
                self.model_person, frame,
                self.config['model']['person'],
                {'person': detection_config['person']}          # 获取人员检测配置
            )
            detections.extend(person_dets)                      # 添加人员检测结果
            has_danger |= person_danger                         # 更新危险标志

            # 处理火焰检测
            fire_dets, fire_danger = self._process_detections(
                self.model_fire, frame,
                self.config['model']['fire'],
                {'fire': detection_config['fire']}      # 获取火焰检测配置
            )
            detections.extend(fire_dets)                            # 添加火焰检测结果
            has_danger |= fire_danger                               # 更新危险标志
            self.dispose = True                                    # 处理完成

            return True, detections, has_danger
        except Exception as e:
            print(f"检测错误: {e}")
            return False, [], False
